import torch
from rich import print
from core import (
    get_dataloader,
    MovieBiLSTM_CRF
)
from utils import conf, data

def debug_comprehensive():
    print('=' * 60)
    print('1. 检查标签映射')
    print('=' * 60)
    # 检查标签文件
    print(f'Labels from file: {data.labels}')
    print(f'Number of labels: {len(data.labels)}')
    # 检查标签是否合理
    if set(data.labels.values()) != set(range(len(data.labels))):
        print('❌ 标签映射有问题!')
        print(f'Expected continuous indices 0-{len(data.labels)-1}')
        print(f'Got: {list(data.labels.values())}')
    else:
        print('✅ 标签映射正常')

    print('\n' + '=' * 60)
    print('2. 检查词汇表')
    print('=' * 60)
    _data, word2id = data.build_vocab()
    print(f'Vocabulary size: {len(word2id)}')
    print(f'Words: {list(word2id.items())[:10]}')
    if 'PAD' not in word2id or word2id['PAD'] != 0:
        print('❌ PAD token 问题!')
    else:
        print('✅ PAD token 正常')

    print('\n' + '=' * 60)
    print('3. 检查数据样本')
    print('=' * 60)
    print(f'Total data samples: {len(_data)}')
    if len(_data) > 0:
        sample = _data[0]
        print(f'Sample text length: {len(sample[0])}')
        print(f'Sample labels length: {len(sample[1])}')
        print(f'Sample text: {sample[0][:20]}')
        print(f'Sample labels: {sample[1][:20]}')
        # 检查标签是否都在预定义标签中
        unique_labels = set(sample[1])
        print(f'Unique labels in sample: {unique_labels}')
        invalid_labels = unique_labels - set(data.labels.keys())
        if invalid_labels:
            print(f'❌ 发现无效标签: {invalid_labels}')
        else:
            print('✅ 所有标签都有效')
    
    print('\n' + '=' * 60)
    print('4. 检查数据加载器')
    print('=' * 60)
    try:
        train_dl, _ = get_dataloader()
        for i, (input_ids, labels, mask) in enumerate(train_dl):
            print(f'Batch {i}:')
            print(f'  Input shape: {input_ids.shape}')
            print(f'  Labels shape: {labels.shape}')
            print(f'  Mask shape: {mask.shape}')
            print(f'  Input range: {input_ids.min().item()} - {input_ids.max().item()}')
            print(f'  Labels range: {labels.min().item()} - {labels.max().item()}')
            print(f'  Mask sum: {mask.sum().item()} / {mask.numel()}')
            # 检查标签是否超出范围
            max_label_id = len(data.labels) - 1
            if labels.max().item() > max_label_id:
                print(f'❌ 标签值超出范围! 最大标签 ID 应该是{max_label_id}, 但发现{labels.max().item()}')
            # 检查padding标签（通常是7或者最大标签+1）
            padding_label = 7
            non_pad_mask = (labels != padding_label)
            valid_labels = labels[non_pad_mask]
            if len(valid_labels) > 0:
                print(f'  有效标签范围: {valid_labels.min().item()} - {valid_labels.max().item()}')
            # 检查第一个样本的详细信息
            if i == 0:
                print(f'  第一个样本详情:')
                seq_len = mask[0].sum().item()
                print(f'    实际序列长度: {seq_len}')
            if i >= 2:
                break
    except Exception as e:
        print(f'❌ 数据加载器错误: {e}')
        import traceback
        traceback.print_exc()

    print('\n' + '=' * 60)
    print('5. 测试模型前向传播')
    print('=' * 60)
    try:
        model = MovieBiLSTM_CRF().to(conf.device)
        print(f'Model device: {next(model.parameters()).device}')
        # 测试一个小批次
        train_dl, _ = get_dataloader()
        input_ids, labels, mask = next(iter(train_dl))
        input_ids = input_ids.to(conf.device)
        labels = labels.to(conf.device)
        mask = mask.to(conf.device)
        print(f'Testing with batch size: {input_ids.shape[0]}')
        # 测试前向传播
        with torch.no_grad():
            output = model(input_ids, mask)
            print(f'Forward pass output length: {len(output)}')
            print(f'First sequence prediction: {output[0][:10] if len(output) > 0 else 'Empty'}')
        # 测试loss计算
        loss = model.log_likelihood(input_ids, labels, mask)
        print(f'Loss value: {loss.item()}')
        if loss.item() > 1000:
            print('❌ Loss 值异常大, 可能的原因:')
            print('   - 标签映射错误')
            print('   - CRF 标签转移矩阵初始化问题')
            print('   - 输入数据范围问题')
        elif loss.item() < 0:
            print('❌ Loss 值为负, 这在 CRF 中不应该发生')
        else:
            print('✅ Loss 值在合理范围内')
    except Exception as e:
        print(f'❌ 模型测试错误: {e}')
        import traceback
        traceback.print_exc()

    print('\n' + '=' * 60)
    print('6. 检查 CRF 配置')
    print('=' * 60)
    print(f'CRF 标签数量: {len(data.labels)}')
    print(f'模型配置的标签数量: {model.tag_size}')
    if hasattr(model.crf, 'start_transitions'):
        print(f'CRF start_transitions shape: {model.crf.start_transitions.shape}')
        print(f'CRF end_transitions shape: {model.crf.end_transitions.shape}')
        print(f'CRF transitions shape: {model.crf.transitions.shape}')
        # 检查转移矩阵是否有异常值
        transitions = model.crf.transitions.data
        print(f'Transitions range: {transitions.min().item():.4f} - {transitions.max().item():.4f}')
        if torch.any(torch.isnan(transitions)):
            print('❌ CRF 转移矩阵包含 NaN 值!')
        elif torch.any(torch.isinf(transitions)):
            print('❌ CRF 转移矩阵包含无穷值!')
        else:
            print('✅ CRF 转移矩阵正常')

if __name__ == '__main__':
    debug_comprehensive()
